/*
 * SPDX-License-Identifier: GPL-2.0
 * Copyright (C) 2025 Skyline <xuemin.zy@hotmail.com>
 * 
 * 基于 mmc-utils (GPLv2) 二次开发，实现 RPMB 协议扩展功能。
 */

#include "rpmb.h"
#include "hmac_sha2.h"
#include "kernel_mmc_ioctl.h"
#include "kernel_mmc_mmc.h"

#define MMC_RSP_NONE	0			/* no response */
#define MMC_RSP_PRESENT	(1 << 0)
#define MMC_RSP_136	(1 << 1)		/* 136 bit response */
#define MMC_RSP_CRC	(1 << 2)		/* expect valid crc */
#define MMC_RSP_BUSY	(1 << 3)		/* card may send busy */
#define MMC_RSP_OPCODE	(1 << 4)		/* response contains opcode */

#define MMC_CMD_AC	(0 << 5)
#define MMC_CMD_ADTC	(1 << 5)
#define MMC_CMD_BC	(2 << 5)

#define MMC_RSP_SPI_S1	(1 << 7)		/* one status byte */
#define MMC_RSP_SPI_BUSY (1 << 10)		/* card may send busy */

#define MMC_RSP_SPI_R1	(MMC_RSP_SPI_S1)
#define MMC_RSP_SPI_R1B	(MMC_RSP_SPI_S1|MMC_RSP_SPI_BUSY)

#define MMC_RSP_R1	(MMC_RSP_PRESENT|MMC_RSP_CRC|MMC_RSP_OPCODE)
#define MMC_RSP_R1B	(MMC_RSP_PRESENT|MMC_RSP_CRC|MMC_RSP_OPCODE|MMC_RSP_BUSY)

#define RPMB_MULTI_CMD_MAX_CMDS 3

#define do_rpmb_op_err(ret) \
{ \
    if (ret != 0) { \
        printz("RPMB ioctl failed ret = %d \r\n", ret); \
        return -1; \
    } \
}

enum rpmb_op_type {
	MMC_RPMB_WRITE_KEY = 0x01,
	MMC_RPMB_READ_CNT  = 0x02,
	MMC_RPMB_WRITE     = 0x03,
	MMC_RPMB_READ      = 0x04,

	/* For internal usage only, do not use it directly */
	MMC_RPMB_READ_RESP = 0x05
};

struct rpmb_frame {
	u_int8_t  stuff[196];
	u_int8_t  key_mac[32];
	u_int8_t  data[RPMB_DAT_BLKSZ];
	u_int8_t  nonce[16];
	u_int32_t write_counter;
	u_int16_t addr;
	u_int16_t block_count;
	u_int16_t result;
	u_int16_t req_resp;
};

static void set_single_cmd(struct mmc_ioc_cmd *p_ioc, unsigned int opcode, int write_flag, unsigned int blocks, unsigned int arg)
{
	p_ioc->opcode = opcode;
	p_ioc->write_flag = write_flag;
	p_ioc->arg = arg;
	p_ioc->blksz = RPMB_BLKSZ;
	p_ioc->blocks = blocks;
	p_ioc->flags = MMC_RSP_SPI_R1 | MMC_RSP_R1 | MMC_CMD_ADTC;
}

static int do_rpmb_op(int fd, const struct rpmb_frame *frame_in, struct rpmb_frame *frame_out, unsigned int blk_cnt)
{
	int err = -1;
	unsigned short rpmb_type = 0;
	struct mmc_ioc_cmd *p_ioc = NULL;
	struct mmc_ioc_multi_cmd *p_mioc = NULL;
	struct rpmb_frame frame_status = {{0}};

	if (!frame_in || !frame_out || !blk_cnt)
    {
        printz("null pointer err \r\n");
        return -1;
    }

	/* prepare arguments for MMC_IOC_MULTI_CMD ioctl */
	p_mioc = (struct mmc_ioc_multi_cmd *)calloc(1, sizeof (struct mmc_ioc_multi_cmd) + RPMB_MULTI_CMD_MAX_CMDS * sizeof (struct mmc_ioc_cmd));
	if (!p_mioc)
    {
        printz("calloc err \r\n");
		return -1;
	}

	rpmb_type = be16toh(frame_in->req_resp);

	switch(rpmb_type)
    {
        case MMC_RPMB_WRITE:
        case MMC_RPMB_WRITE_KEY:
            if (blk_cnt != 1)
            {
                err = -1;
                goto out;
            }

            p_mioc->num_of_cmds = 3;

            /* Write request */
            p_ioc = &p_mioc->cmds[0];
            set_single_cmd(p_ioc, MMC_WRITE_MULTIPLE_BLOCK, (1 << 31) | 1, blk_cnt, 0);
            mmc_ioc_cmd_set_data((*p_ioc), frame_in);

            /* Result request */
            p_ioc = &p_mioc->cmds[1];
            frame_status.req_resp = htobe16(MMC_RPMB_READ_RESP);
            set_single_cmd(p_ioc, MMC_WRITE_MULTIPLE_BLOCK, 1, 1, 0);
            mmc_ioc_cmd_set_data((*p_ioc), &frame_status);

            /* Get response */
            p_ioc = &p_mioc->cmds[2];
            set_single_cmd(p_ioc, MMC_READ_MULTIPLE_BLOCK, 0, 1, 0);
            mmc_ioc_cmd_set_data((*p_ioc), frame_out);

            break;
        case MMC_RPMB_READ_CNT:
            if (blk_cnt != 1)
            {
                err = -1;
                goto out;
            }
            /* fall through */

        case MMC_RPMB_READ:
            p_mioc->num_of_cmds = 2;

            /* Read request */
            p_ioc = &p_mioc->cmds[0];
            set_single_cmd(p_ioc, MMC_WRITE_MULTIPLE_BLOCK, 1, 1, 0);
            mmc_ioc_cmd_set_data((*p_ioc), frame_in);

            /* Get response */
            p_ioc = &p_mioc->cmds[1];
            set_single_cmd(p_ioc, MMC_READ_MULTIPLE_BLOCK, 0, blk_cnt, 0);
            mmc_ioc_cmd_set_data((*p_ioc), frame_out);

            break;
        default:
            err = -1;
            goto out;
	}

    printz("rpmb_type = %d \r\n", rpmb_type);

	err = ioctl(fd, MMC_IOC_MULTI_CMD, p_mioc);
out:
	free(p_mioc);
	return err;
}

int do_rpmb_write_key(int dev_fd, char *p_rpmb_key, unsigned int key_size)
{
    int ret = -1;
    struct rpmb_frame frame_in = {0};
    struct rpmb_frame frame_out = {0};

    if (dev_fd <= 2)
    {
        printz("dev_fd[%d] err \r\n", dev_fd);
    }

    if (NULL == p_rpmb_key)
    {
        printz("null pointer err \r\n");
        return -1;
    }

    if (key_size != sizeof(frame_in.key_mac))
    {
        printz("key_size[%d] err \r\n", key_size);
        return -1;
    }

    frame_in.req_resp = htobe16(MMC_RPMB_WRITE_KEY);
    memcpy(frame_in.key_mac, p_rpmb_key, key_size);

	/* Execute RPMB op */
	do_rpmb_op_err(ret = do_rpmb_op(dev_fd, &frame_in, &frame_out, 1));

	/* Check RPMB response */
	if (frame_out.result != 0)
    {
		printz("RPMB operation failed, retcode 0x%04x \r\n", be16toh(frame_out.result));
		return -1;
	}

	return 0;
}

int do_rpmb_read_counter(int dev_fd, unsigned int *p_counter)
{
	int ret;
    struct rpmb_frame frame_in = {0};
    struct rpmb_frame frame_out = {0};

    if (dev_fd <= 2)
    {
        printz("dev_fd[%d] err \r\n", dev_fd);
    }

    if (NULL == p_counter)
    {
        printz("null pointer err \r\n");
        return -1;
    }

    frame_in.req_resp = htobe16(MMC_RPMB_READ_CNT);

	/* Execute RPMB op */
	do_rpmb_op_err(ret = do_rpmb_op(dev_fd, &frame_in, &frame_out, 1));

	/* Check RPMB response */
	if (frame_out.result != 0)
    {
		*p_counter = 0;
        printz("operation failed, retcode = 0x%04x \r\n", frame_out.result);
		return be16toh(frame_out.result);
	}

	*p_counter = be32toh(frame_out.write_counter);

	return 0;
}

int do_rpmb_read_block(int dev_fd, unsigned short addr, unsigned int blocks_cnt, char *p_out, unsigned int outbuf_size, char *p_rpmb_key, unsigned int rpmb_keysize)
{
	int i = 0;
    int ret = -1;
    unsigned char mac[32];
	hmac_sha256_ctx ctx = {{0}};
    struct rpmb_frame frame_in = {{0}};
    struct rpmb_frame *p_frame_out = NULL;
    struct rpmb_frame *p_frame_tmp = NULL;

    if (dev_fd <= 2)
    {
        printz("dev_fd[%d] err \r\n", dev_fd);
    }

    if (NULL == p_out)
    {
        printz("null pointer err \r\n");
        return -1;
    }

    if (!blocks_cnt)
    {
		printz("please, specify valid blocks count number \r\n");
		return -1;
	}

    if (RPMB_DAT_BLKSZ * blocks_cnt > outbuf_size)
    {
        printz("output buf size err, [%d] > [%d] \r\n", RPMB_DAT_BLKSZ * blocks_cnt, outbuf_size);
        return -1;
    }

	frame_in.addr = htobe16(addr);
    frame_in.req_resp = htobe16(MMC_RPMB_READ);
	p_frame_out = calloc(sizeof(struct rpmb_frame), blocks_cnt);
	if (!p_frame_out)
    {
		printz("can't allocate memory for RPMB outer frames \r\n");
		return -1;
	}

	/* Execute RPMB op */
	ret = do_rpmb_op(dev_fd, &frame_in, p_frame_out, blocks_cnt);
    if (ret)
    {
        printz("do_rpmb_op err \r\n");
        free(p_frame_out);
		return -1;
    }

	/* Check RPMB response */
	if (p_frame_out[blocks_cnt - 1].result != 0)
    {
		printz("RPMB operation failed, retcode 0x%04x \r\n", be16toh(p_frame_out[blocks_cnt - 1].result));
        free(p_frame_out);
		return -1;
	}

	/* Do we have to verify data against key? */
	if (NULL != p_rpmb_key)
    {
        if (RPMB_KEY_STDSIZE != rpmb_keysize)
        {
            printz("key_size[%d] err \r\n", rpmb_keysize);
        }

		hmac_sha256_init(&ctx, p_rpmb_key, RPMB_KEY_STDSIZE);
		for (i = 0; i < blocks_cnt; i++)
        {
			p_frame_tmp = &p_frame_out[i];
			hmac_sha256_update(&ctx, p_frame_tmp->data, sizeof(*p_frame_tmp) - offsetof(struct rpmb_frame, data));
		}

		hmac_sha256_final(&ctx, mac, sizeof(mac));

		/* Impossible */
		assert(p_frame_tmp);

		/* Compare calculated MAC and MAC from last frame */
		if (memcmp(mac, p_frame_tmp->key_mac, sizeof(mac)))
        {
			printz("RPMB MAC missmatch \r\n");
            free(p_frame_out);
			return -1;
		}
	}

	/* Write data */
    for (i = 0; i < blocks_cnt; i++)
    {
        p_frame_tmp = &p_frame_out[i];
        memcpy(p_out + sizeof(p_frame_tmp->data) * i, p_frame_tmp->data, sizeof(p_frame_tmp->data));
    }

	free(p_frame_out);

	return ret;
}

int do_rpmb_write_block(int dev_fd, unsigned short addr, char *p_in, unsigned int in_len, char *p_rpmb_key, unsigned int rpmb_keysize)
{
	int ret = -1;
    unsigned int i = 0;
	unsigned int cnt = 0;
    struct rpmb_frame frame_in = {{0}};
	struct rpmb_frame frame_out = {{0}};

    if (dev_fd <= 2)
    {
        printz("dev_fd[%d] err \r\n", dev_fd);
    }

    if (NULL == p_in)
    {
        printz("null pointer err \r\n");
        return -1;
    }

    if (RPMB_KEY_STDSIZE != rpmb_keysize)
    {
        printz("key_size[%d] err \r\n", rpmb_keysize);
        return -1;
    }

    if (in_len > sizeof(frame_in.data))
    {
        printz("in_len[%d] err \r\n", in_len);
        return -1;
    }

	ret = do_rpmb_read_counter(dev_fd, &cnt);
	if (ret != 0)
    {
		printz("RPMB read counter operation failed, retcode 0x%04x\n", ret);
		return -1;
	}

    frame_in.req_resp = htobe16(MMC_RPMB_WRITE);
    frame_in.block_count = htobe16(1);
    frame_in.write_counter = htobe32(cnt);
    frame_in.addr = htobe16(addr);
    memcpy(frame_in.data, p_in, sizeof(frame_in.data));
    memset(frame_in.data + in_len, 'F', sizeof(frame_in.data) - in_len);

	/* Calculate HMAC SHA256 */
	hmac_sha256(p_rpmb_key, rpmb_keysize,
		        frame_in.data, sizeof(struct rpmb_frame) - offsetof(struct rpmb_frame, data),
		        frame_in.key_mac, sizeof(frame_in.key_mac));

	/* Execute RPMB op */
	ret = do_rpmb_op(dev_fd, &frame_in, &frame_out, 1);
    if (ret)
    {
        printz("do_rpmb_op err \r\n");
    }
    else
    {
        /* Check RPMB response */
        if (frame_out.result != 0)
        {
            printz("RPMB operation failed, retcode 0x%04x\n", be16toh(frame_out.result));
            return -1;
        }
    }

	return ret;
}

